import sys
import os
import random
import tensorflow as tf
import math
print("Tensorflow version " + tf.__version__)
tf.set_random_seed(0)
import numpy as np
import cv2
import next_batch

X = tf.placeholder(tf.float32, [None, 28, 28, 1])
Y_= tf.placeholder(tf.float32, [None, 10])
keep_prob = tf.placeholder(tf.float32)

K = 12  # 1 convolutional layer output depth
L = 24  # 2 convolutional layer output depth
M = 36  # 3 convolutional layer
N = 256 # 4 fully connected layer
C = 10  # 5 fully connection layer
W1 = tf.Variable(tf.truncated_normal([3, 3, 4, K], stddev=0.1)) 
B1 = tf.Variable(tf.ones([K]) / 10)
W2 = tf.Variable(tf.truncated_normal([3, 3, K, L], stddev=0.1))
B2 = tf.Variable(tf.ones([L]) / 10)
W3 = tf.Variable(tf.truncated_normal([3, 3, L, M], stddev=0.1))
B3 = tf.Variable(tf.ones([M]) / 10)
W4 = tf.Variable(tf.truncated_normal([7* 7 *M, N], stddev=0.1))
B4 = tf.Variable(tf.ones([N]) / 10)
W5 = tf.Variable(tf.truncated_normal([N,      10], stddev=0.1))
B5 = tf.Variable(tf.ones([C]) / 10)
# add the position sensitive factor, after this op, the image size 14 *14
s2b= tf.space_to_depth(X, block_size = 2)
#
C1 = tf.nn.conv2d(s2b,  W1, strides=[1, 1, 1, 1], padding='SAME')
Y1 = tf.nn.relu(C1 + B1)
#
C2 = tf.nn.conv2d(Y1, W2, strides=[1, 1, 1, 1], padding='SAME')
Y2 = tf.nn.relu(C2 + B2)
# output is 7 x 7
C3 = tf.nn.conv2d(Y2, W3, strides=[1, 1, 1, 1], padding='SAME')
Y3 = tf.nn.relu(C3 + B3)
P3 = tf.nn.max_pool(Y3, ksize = [1,2,2,1], strides = [1,2,2,1], padding='SAME') 
# full connection
YY = tf.reshape(P3, shape=[-1, 7 * 7 * M])
Y4 = tf.nn.relu(tf.matmul(YY, W4) + B4)
D1 = tf.nn.dropout(Y4, keep_prob)
# full connection
Ylogits = tf.matmul(D1, W5) + B5
Y = tf.nn.softmax(Ylogits)
# cross entropy
cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits=Ylogits, labels=Y_)
cross_entropy = tf.reduce_mean(cross_entropy) * 100
# accuracy of the trained model, between 0 (worst) and 1 (best)
correct_prediction = tf.equal(tf.argmax(Y, 1), tf.argmax(Y_, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
# training step, the learning rate is a placeholder
train_step = tf.train.AdamOptimizer(1e-5).minimize(cross_entropy)
# init
init = tf.global_variables_initializer()
# save model
Saver = tf.train.Saver(max_to_keep = 1)  # defaults to saving all variables
ModelSavePath = './space2depth_storage/'

def train():
	with tf.Session() as sess:
		sess.run(init)
		ckpt = tf.train.get_checkpoint_state(ModelSavePath)
		if ckpt and ckpt.model_checkpoint_path:
			Saver.restore(sess, ckpt.model_checkpoint_path)
			print '>> infor: restore model ok'
		else:
			print '>> infor: model so far'
		# prepare dataset
		dataset = next_batch.batchData()
		iterNum = 10000000
		batchsize = 500
		for i in xrange(iterNum):
			batch_X, batch_Y = dataset.next_batch_train(batchsize)
			a, c, s = sess.run([accuracy, cross_entropy, train_step], {X: batch_X, Y_: batch_Y, keep_prob: 0.5})
			if (i + 1) % 50 == 0:
				batch_X, batch_Y = dataset.all_test()
				test_accu, test_Y = sess.run([accuracy, Y], {X: batch_X, Y_: batch_Y, keep_prob: 1.0})
				fg = 0
				if (i + 1) % 1000 == 0:
					# save tensorflow model
					Saver.save(sess, ModelSavePath + "model", global_step = (i + 1))
					# save W and B into npz file
					np.savez(
						ModelSavePath + 'model.npz', \
						W1 = sess.run(W1), \
						B1 = sess.run(B1), \
						W2 = sess.run(W2), \
						B2 = sess.run(B2), \
						W3 = sess.run(W3), \
						B3 = sess.run(B3), \
						W4 = sess.run(W4), \
						B4 = sess.run(B4), \
						W5 = sess.run(W5), \
						B5 = sess.run(B5))
					fg = 1
				radIdx = random.randint(0, len(batch_Y))
				if fg == 1:
					print iterNum, 	'%6d'%(i + 1), '%9.4f(cross entropy) '%c, \
					'%9.4f(training accuracy) '%a, '%9.4f(test accuracy) '%test_accu, \
					np.argmax(batch_Y, 1)[radIdx:radIdx + 5], np.argmax(test_Y, 1)[radIdx:radIdx + 5], \
					'save'
				else:
					print iterNum, 	'%6d'%(i + 1), '%9.4f(cross entropy) '%c, \
					'%9.4f(training accuracy) '%a, '%9.4f(test accuracy) '%test_accu, \
					np.argmax(batch_Y, 1)[radIdx:radIdx + 5], np.argmax(test_Y, 1)[radIdx:radIdx + 5]

def test():
	batch_X = np.zeros([27, 28, 28, 1])
	for i in range(0, 9):
		for j in range(0, 3):
			curF = '../test_data1/z' + '%d'%i + '%d'%j + '.jpg'
			batch_X[i*3 + j] = np.expand_dims(cv2.imread(curF, cv2.IMREAD_GRAYSCALE).astype('float32')/ 255.0, -1)
			batch_Y = [5, 6, 7, 
			           0, 0, 1, 
			           0, 1, 2, 
			           4, 5, 6, 
			           6, 9, 8, 
			           5, 6, 7, 
			           9, 8, 0, 
			           0, 1, 2,
			           4, 5, 6]
	with tf.Session() as sess:
		sess.run(init)
		ckpt = tf.train.get_checkpoint_state(ModelSavePath)
		if ckpt and ckpt.model_checkpoint_path:
			Saver.restore(sess, ckpt.model_checkpoint_path)
			print '>> infor: restore model ok'
		else:
			print '>> infor: model so far'
		oy = sess.run(Y, {X: batch_X, keep_prob:1.0})
		for i in range(0,24):
			if batch_Y[i] == np.argmax(oy, 1)[i]:
				print batch_Y[i], np.argmax(oy, 1)[i]
			else:
				print batch_Y[i], np.argmax(oy, 1)[i], 'X'
#		print oy

if __name__ == "__main__":
    if (len(sys.argv) == 2):
    	if sys.argv[1] == '0':
    		train()
    	else:
    		test()
    else:
        print '>> error. input parameters'


